// qwen-llm.ts
import { BaseLLM } from "langchain/llms/base";
import { CallbackManagerForLLMRun } from "langchain/callbacks";
import { AxiosRequestConfig } from "axios";
import axios from "axios";
import { BaseLLMParams } from "langchain/llms/base";
import { LLMResult } from "langchain/schema";

interface QwenCallArgs {
  model: string; // e.g. qwen-max, qwen-plus, qwen-turbo
  temperature?: number;
  max_tokens?: number;
}

export class QwenLLM extends BaseLLM {
  apiUrl = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation";

  model: string;
  temperature: number;
  maxTokens: number;
  apiKey: string;

  constructor(
    apiKey: string,
    config: QwenCallArgs = { model: "qwen-max" }
  ) {
    super({});
    this.apiKey = apiKey;
    this.model = config.model || "qwen-max";
    this.temperature = config.temperature || 0.7;
    this.maxTokens = config.max_tokens || 512;
  }

  // 实现抽象方法 _generate
  async _generate(
    prompts: string[],
    options: this["ParsedCallOptions"],
    runManager?: CallbackManagerForLLMRun
  ): Promise<LLMResult> {
    const generations = [];
    for (const prompt of prompts) {
      const text = await this._call(prompt, options, runManager);
      generations.push([{ text }]);
    }
    return { generations };
  }

  // 核心调用方法
  async _call(
    prompt: string,
    options?: this["ParsedCallOptions"],
    runManager?: CallbackManagerForLLMRun
  ): Promise<string> {
    const headers: AxiosRequestConfig["headers"] = {
      "Authorization": `Bearer ${this.apiKey}`,
      "Content-Type": "application/json",
    };

    const data = {
      model: this.model,
      input: {
        prompt,
      },
      parameters: {
        temperature: this.temperature,
        max_tokens: this.maxTokens,
      },
    };

    try {
      const response = await axios.post(this.apiUrl, data, { headers });
      return response.data.output.text;
    } catch (error: any) {
      console.error("Qwen API Error:", error.response?.data || error.message);
      throw new Error(
        `Qwen call failed: ${error.response?.data?.message || error.message}`
      );
    }
  }

  _llmType(): string {
    return "qwen";
  }
}